import qianfan
from zhipuai import ZhipuAI
import os

class LLM:
    def __init__(self):
        self.zhipuai_key = "06f08c8b4a634719f10ee63a802c1cc2.fCWxL1JzrQDVc0RL"
        os.environ["QIANFAN_ACCESS_KEY"] = "217a3a9a6136459680bdde834e246f69"
        os.environ["QIANFAN_SECRET_KEY"] = "d1b9b7624e05402f8a614605bb6e1107"

    async def ernie(self, input_text, history, model):
        # ERNIE-4.0-Turbo-8K：60 8K
        # ERNIE-Speed-128K：免费 128K
        # ERNIE-Tiny-8K：免费 8K

        chat_comp = qianfan.ChatCompletion()
        if history != []:
            del history[0]
        messages = history
        messages.append({"role": "user", "content": f"{input_text}"})
        resp = chat_comp.do(
            model=model,
            messages=messages
        )
        output = resp["body"]['result']
        return output


    async def chatglm(self, input_text, history, model="glm-4-flash"):
        """
        # chatglm 可用模型（单位 元/百万tokens）：
        # glm-4：100 128K
        # glm-4-air：1 128K
        # glm-4-flash：0.1 128K
        """
        client = ZhipuAI(api_key="06f08c8b4a634719f10ee63a802c1cc2.fCWxL1JzrQDVc0RL")

        messages = history
        messages.append({"role": "user", "content": f"{input_text}"})
        response = client.chat.completions.create(
            model=model,  # 填写需要调用的模型名称
            messages=messages,
        )
        output = response.choices[0].message.content
        return output